#include <sys/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <netdb.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <arpa/inet.h> 
#include "util.h"
#include <inttypes.h>
#include <time.h>

/*******************************************************
  这个是Client进程，用来发起请求。这里做一个简化，省去登
  陆的步骤，以及Master Key的生成步骤。这里直接假定Master
  Key是0x6162636465666768。
********************************************************/
#define CLIENT_ID "Bob"
#define K_CLIENT 0x6162636465666768

int main() {
    int sockfd = 0;
    struct sockaddr_in serv_addr; 

    /* 开始与AS连接 */
    if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
        printf("Error : Could not create socket \n");
        return 1;
    }

    memset(&serv_addr, 0, sizeof(serv_addr)); 

    serv_addr.sin_family = AF_INET;
    serv_addr.sin_port = htons(11111); 

    if (inet_pton(AF_INET, "127.0.0.1", &serv_addr.sin_addr) <= 0) {
        printf("inet_pton error occured\n");
        return 1;
    }

    int res = connect(sockfd, (struct sockaddr *)&serv_addr, sizeof(serv_addr));
    if (res < 0) {
       printf("errno %d: %s\n", errno, strerror(errno));
       return 1;
    }
    printf("Connected to AS!\n");

    /* 与AS建立连接后，首先就要明文发送自己的Client ID给AS */
    multi_msg plain1 = new_multi_msg();
    mm_build(&plain1, "s", CLIENT_ID);
    mm_sock_write(sockfd, &plain1);
    mm_free(&plain1);

    /* 收到AB两则消息 */
    multi_msg message_ab = new_multi_msg();
    mm_sock_read(&message_ab, sockfd);
    if (strcmp((char*)message_ab.msgs[0].data, "success") != 0) {
        printf("Client ID not exists\n");
        return 1;
    }
    multi_msg message_a = new_multi_msg();
    mm_decrypt_loads(&message_a, message_ab.msgs + 1, K_CLIENT);
    uint64_t k_client_tgs = 0;
    for (int i = 0; i < 8; ++i) {
        k_client_tgs |= ((uint64_t)message_a.msgs[0].data[i] << (8 * i));
    }
    printf("Decrypt and get the new k_client_tgs: %016" PRIx64 "\n", k_client_tgs);

    /* 注意到Client是无法解密B消息(TGT)的 */
    octet_string eos_b = new_octet_string();
    eos_b.size = message_ab.msgs[2].size;
    eos_b.data = (uint8_t*)malloc(eos_b.size);
    memcpy(eos_b.data, message_ab.msgs[2].data, eos_b.size);
    mm_free(&message_ab);
    printf("Get TGT: ");
    for (size_t i = 0; i < eos_b.size; ++i) {
        printf("%02X", eos_b.data[i]);
    }
    printf("\n\n");
    
    close(sockfd);




    /* 开始与TGS连接 */
    if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
        printf("Error : Could not create socket \n");
        return 1;
    }

    serv_addr.sin_family = AF_INET;
    serv_addr.sin_port = htons(22222); 

    if (inet_pton(AF_INET, "127.0.0.1", &serv_addr.sin_addr) <= 0) {
        printf("inet_pton error occured\n");
        return 1;
    } 

    res = connect(sockfd, (struct sockaddr *)&serv_addr, sizeof(serv_addr));
    if (res < 0) {
       printf("errno %d: %s\n", errno, strerror(errno));
       return 1;
    } 
    printf("Connected to TGS!\n");

    /* 发送CD两则消息 */
    multi_msg message_c = new_multi_msg();
    mm_build(&message_c, "so", "toupper", &eos_b);
    octet_string os_c = new_octet_string();
    mm_dumps(&os_c, &message_c);
    mm_free(&message_c);

    multi_msg message_d = new_multi_msg();
    mm_build(&message_d, "sd", CLIENT_ID, time(NULL));
    octet_string os_d = new_octet_string();
    mm_dumps_encrypt(&os_d, &message_d, k_client_tgs);
    mm_free(&message_d);

    multi_msg message_cd = new_multi_msg();
    mm_build(&message_cd, "oo", &os_c, &os_d);
    os_free(&os_c);
    os_free(&os_d);
    mm_sock_write(sockfd, &message_cd);
    mm_free(&message_cd);

    /* 接受EF两则消息 */
    multi_msg message_ef = new_multi_msg();
    mm_sock_read(&message_ef, sockfd);
    if (strcmp((char*)message_ef.msgs[0].data, "success") != 0) {
        printf("Authentication failed\n");
        return 1;
    }
    /* 从F中提取k_client_ss */
    multi_msg message_f = new_multi_msg();
    mm_decrypt_loads(&message_f, message_ef.msgs + 2, k_client_tgs);
    uint64_t k_client_ss = 0;
    for (int i = 0; i < 8; ++i) {
        k_client_ss |= ((uint64_t)message_f.msgs[0].data[i] << (i * 8));
    }
    mm_free(&message_f);
    printf("Decrypt and get the new k_client_ss: %016" PRIx64 "\n", k_client_ss);

    multi_msg message_e = new_multi_msg();
    mm_loads(&message_e, message_ef.msgs + 1);
    mm_free(&message_ef);
    printf("Get ST: ");
    for (int i = 0; i < message_e.msgs[1].size; ++i) {
        printf("%02X", message_e.msgs[1].data[i]);
    }

    printf("\n\n");
    close(sockfd);





    /* 开始与SS连接 */
    if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
        printf("Error : Could not create socket \n");
        return 1;
    }

    serv_addr.sin_family = AF_INET;
    serv_addr.sin_port = htons(33333); 

    if (inet_pton(AF_INET, "127.0.0.1", &serv_addr.sin_addr) <= 0) {
        printf("inet_pton error occured\n");
        return 1;
    }

    res = connect(sockfd, (struct sockaddr *)&serv_addr, sizeof(serv_addr));
    if (res < 0) {
       printf("errno %d: %s\n", errno, strerror(errno));
       return 1;
    }
    printf("Connected to SS!\n");

    /* 发送EG两则消息 */
    multi_msg message_g = new_multi_msg();
    uint32_t stamp = time(NULL);
    mm_build(&message_g, "sd", CLIENT_ID, stamp);
    octet_string eos_g = new_octet_string();
    mm_dumps_encrypt(&eos_g, &message_g, k_client_ss);
    mm_free(&message_g);

    octet_string os_e = new_octet_string();
    mm_dumps(&os_e, &message_e);
    mm_free(&message_e);

    multi_msg message_eg = new_multi_msg();
    mm_build(&message_eg, "oo", &os_e, &eos_g);
    mm_sock_write(sockfd, &message_eg);


    /* 接受H消息，检验timestamp是否被正确+1 */
    multi_msg message_final = new_multi_msg();
    mm_sock_read(&message_final, sockfd);
    multi_msg message_h = new_multi_msg();
    mm_decrypt_loads(&message_h, message_final.msgs + 1, k_client_ss);
    uint32_t stamp_plus_1 = 0;
    for (int i = 0; i < 4; ++i) {
        stamp_plus_1 |= ((uint32_t)message_h.msgs[1].data[i] << (i * 8));
    }
    if (stamp_plus_1 != stamp + 1 || strcmp((char*)message_h.msgs[0].data, CLIENT_ID) != 0) {
        printf("Authentication failed.\n");
        return 1;
    }
    printf("Authentication success. Start using service provided by SS from now on.\n\n");




    /* 开始使用SS提供的转大写服务 */
    const char *test_data[] = {
        "The quick brown fox jumps over the lazy dog.",
        "I am Zhang3 and I study in SYSU",
        "LBWNB!!",
        "7355608",
    };
    size_t num_data = sizeof(test_data) / sizeof(test_data[0]);
    for (size_t i = 0; i < num_data; ++i) {
        multi_msg message_test = new_multi_msg();
        if (i == num_data - 1) {
            mm_build(&message_test, "ss", test_data[i], "end");
        }
        else {
            mm_build(&message_test, "s", test_data[i]);
        }
        octet_string eos_test = new_octet_string();
        mm_dumps_encrypt(&eos_test, &message_test, k_client_ss);
        mm_free(&message_test);
        multi_msg message_cipher = new_multi_msg();
        mm_build(&message_cipher, "o", &eos_test);
        os_free(&eos_test);
        mm_sock_write(sockfd, &message_cipher);
        mm_free(&message_cipher);


        mm_sock_read(&message_cipher, sockfd);
        mm_decrypt_loads(&message_test, message_cipher.msgs + 0, k_client_ss);
        mm_free(&message_cipher);
        printf("before toupper: %s\n", test_data[i]);
        printf(" after toupper: %s\n", message_test.msgs[0].data);
        mm_free(&message_test);
    }

    close(sockfd);

    return 0;
}
